{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72922059-0f22-467c-abe3-9bc23139adc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import random\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "169885d7-bfa5-4297-a8c3-e9f559d73b14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow version: 2.5.0\n",
      "Detected GPUs: 1\n"
     ]
    }
   ],
   "source": [
    "print(\"TensorFlow version:\", tf.__version__)\n",
    "\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "print(\"Detected GPUs:\", len(gpus))\n",
    "for gpu in gpus:\n",
    "    tf.config.experimental.set_memory_growth(gpu, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e37a2821-ab98-4797-a37b-93b4e7df80ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.initializers import RandomNormal\n",
    "import tensorflow.keras.backend as K\n",
    "\n",
    "#Modified depth_to_space shuffle order for easier shader generation\n",
    "class DepthToSpace2(tf.keras.layers.Layer):\n",
    "    def __init__(self, input_depth, **kwargs):\n",
    "        super(DepthToSpace2, self).__init__(**kwargs)\n",
    "        self.input_depth = input_depth\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        super(DepthToSpace2, self).build(input_shape)\n",
    "\n",
    "    def call(self, x):\n",
    "        x = tf.split(x, (self.input_depth // 4), axis=-1)\n",
    "        return tf.concat([tf.nn.depth_to_space(xx, 2) for xx in x], axis=-1)\n",
    "\n",
    "#SR model that does not change image size\n",
    "def SR1Model(input_texture=\"MAIN\", input_depth=3, highway_depth=4, block_depth=4, init='he_normal', init_last = RandomNormal(mean=0.0, stddev=0.001)):\n",
    "\n",
    "    input_shape = [None, None, input_depth]\n",
    "    #Add \".MAIN\" in layer name as flag for shader generation, this makes the input act as the MAIN texture\n",
    "    input_lr = tf.keras.layers.Input(shape=input_shape, name=\"input.\" + input_texture)\n",
    "    \n",
    "    depth_list = []\n",
    "    \n",
    "    x = input_lr\n",
    "    for i in range(block_depth):\n",
    "        x = tf.keras.layers.Conv2D(highway_depth, (3, 3), padding='same', kernel_initializer=init)(x)\n",
    "        x = tf.nn.crelu(x)\n",
    "        depth_list.append(x)\n",
    "\n",
    "    x = tf.keras.layers.Concatenate(axis=-1)(depth_list)\n",
    "    \n",
    "    #Add \"lastresid\" in layer name as flag for shader generation, this allows the shader to combine the convolution with the residual add as one layer for faster performance\n",
    "    #Add \".MAIN\" in layer name to make the layer save to the MAIN texture\n",
    "    x = tf.keras.layers.Conv2D(input_depth, (1, 1), padding='same', kernel_initializer=init_last, name=\"conv2d_lastresid.\" + input_texture)(x)\n",
    "    \n",
    "    #Add \".ignore\" in layer name as flag for shader generation, this will ignore the layer, as the residual will be added by the previous \"lastresid\" layer\n",
    "    x = tf.keras.layers.Add(name=\"add.ignore.\" + input_texture)([x, input_lr])\n",
    "\n",
    "    model = tf.keras.models.Model(input_lr, x)\n",
    "\n",
    "    return model\n",
    "\n",
    "#SR model that doubles image size\n",
    "def SR2Model(input_texture=\"MAIN\", input_depth=3, highway_depth=4, block_depth=4, init='he_normal', init_last = RandomNormal(mean=0.0, stddev=0.001)):\n",
    "\n",
    "    input_shape = [None, None, input_depth]\n",
    "    #Add \".MAIN\" in layer name as flag for shader generation, this makes the input act as the MAIN texture\n",
    "    input_lr = tf.keras.layers.Input(shape=input_shape, name=\"input.\" + input_texture)\n",
    "    input_lr2 = tf.keras.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(input_lr)\n",
    "    \n",
    "    depth_list = []\n",
    "    \n",
    "    x = input_lr\n",
    "    for i in range(block_depth):\n",
    "        x = tf.keras.layers.Conv2D(highway_depth, (3, 3), padding='same', kernel_initializer=init)(x)\n",
    "        x = tf.nn.crelu(x)\n",
    "        depth_list.append(x)\n",
    "\n",
    "    x = tf.keras.layers.Concatenate(axis=-1)(depth_list)\n",
    "    x = tf.keras.layers.Conv2D(4*input_depth, (1, 1), padding='same', kernel_initializer=init_last)(x)\n",
    "    \n",
    "    #Add \"lastresid\" in layer name as flag for shader generation, this allows the shader to combine the convolution with the residual add as one layer for faster performance\n",
    "    #Add \".MAIN\" in layer name to make the layer save to the MAIN texture\n",
    "    x = DepthToSpace2(4*input_depth, name=\"depth_to_space2_lastresid.\" + input_texture)(x)\n",
    "    \n",
    "    #Add \".ignore\" in layer name as flag for shader generation, this will ignore the layer, as the residual will be added by the previous \"lastresid\" layer\n",
    "    x = tf.keras.layers.Add(name=\"add.ignore.\" + input_texture)([x, input_lr2])\n",
    "\n",
    "    model = tf.keras.models.Model(input_lr, x)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0b7d177e-7e9e-461b-83d4-60146f323250",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "Layer (type)                                     Output Shape                     Param #           Connected to                                      \n",
      "======================================================================================================================================================\n",
      "input.MAIN (InputLayer)                          [(None, None, None, 3)]          0                                                                   \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d (Conv2D)                                  (None, None, None, 4)            112               input.MAIN[0][0]                                  \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu (TFOpLambda)               (None, None, None, 8)            0                 conv2d[0][0]                                      \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_1 (Conv2D)                                (None, None, None, 4)            292               tf.compat.v1.nn.crelu[0][0]                       \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu_1 (TFOpLambda)             (None, None, None, 8)            0                 conv2d_1[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_2 (Conv2D)                                (None, None, None, 4)            292               tf.compat.v1.nn.crelu_1[0][0]                     \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu_2 (TFOpLambda)             (None, None, None, 8)            0                 conv2d_2[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_3 (Conv2D)                                (None, None, None, 4)            292               tf.compat.v1.nn.crelu_2[0][0]                     \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu_3 (TFOpLambda)             (None, None, None, 8)            0                 conv2d_3[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_4 (Conv2D)                                (None, None, None, 4)            292               tf.compat.v1.nn.crelu_3[0][0]                     \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu_4 (TFOpLambda)             (None, None, None, 8)            0                 conv2d_4[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_5 (Conv2D)                                (None, None, None, 4)            292               tf.compat.v1.nn.crelu_4[0][0]                     \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu_5 (TFOpLambda)             (None, None, None, 8)            0                 conv2d_5[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_6 (Conv2D)                                (None, None, None, 4)            292               tf.compat.v1.nn.crelu_5[0][0]                     \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "tf.compat.v1.nn.crelu_6 (TFOpLambda)             (None, None, None, 8)            0                 conv2d_6[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "concatenate (Concatenate)                        (None, None, None, 56)           0                 tf.compat.v1.nn.crelu[0][0]                       \n",
      "                                                                                                    tf.compat.v1.nn.crelu_1[0][0]                     \n",
      "                                                                                                    tf.compat.v1.nn.crelu_2[0][0]                     \n",
      "                                                                                                    tf.compat.v1.nn.crelu_3[0][0]                     \n",
      "                                                                                                    tf.compat.v1.nn.crelu_4[0][0]                     \n",
      "                                                                                                    tf.compat.v1.nn.crelu_5[0][0]                     \n",
      "                                                                                                    tf.compat.v1.nn.crelu_6[0][0]                     \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "conv2d_7 (Conv2D)                                (None, None, None, 12)           684               concatenate[0][0]                                 \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "depth_to_space2_lastresid.MAIN (DepthToSpace2)   (None, None, None, 3)            0                 conv2d_7[0][0]                                    \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "up_sampling2d (UpSampling2D)                     (None, None, None, 3)            0                 input.MAIN[0][0]                                  \n",
      "______________________________________________________________________________________________________________________________________________________\n",
      "add.ignore.MAIN (Add)                            (None, None, None, 3)            0                 depth_to_space2_lastresid.MAIN[0][0]              \n",
      "                                                                                                    up_sampling2d[0][0]                               \n",
      "======================================================================================================================================================\n",
      "Total params: 2,548\n",
      "Trainable params: 2,548\n",
      "Non-trainable params: 0\n",
      "______________________________________________________________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "K.reset_uids()\n",
    "model = SR2Model(input_texture=\"MAIN\", input_depth=3, highway_depth=4, block_depth=7)\n",
    "model.summary(line_length=150)\n",
    "model.load_weights(\"model-checkpoint.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c7d1a6fa-eb73-4dfb-b50a-ee71e1e6481f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from shaderutils import gen_shader\n",
    "gen_shader(model, hook=\"MAIN\", file=\"Upscale_Shader.glsl\", desc=\"Upscale\", when=\"OUTPUT.w MAIN.w / 1.200 > OUTPUT.h MAIN.h / 1.200 > *\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
